#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高斯数据库表字段添加脚本（最终版）
功能：读取columns.txt文件，生成完整的ALTER TABLE语句
特点：
1. 支持以数字开头的字段名（使用双引号）
2. 完整的错误处理和验证
3. 生成可直接执行的SQL脚本
4. 包含详细的执行说明
作者：AI Assistant
"""

import os
import re
from collections import defaultdict

def map_data_type_final(original_type):
    """
    最终版数据类型映射
    
    Args:
        original_type (str): 原始数据类型
        
    Returns:
        str: 映射后的数据类型
    """
    type_mapping = {
        'bigint': 'BIGINT',
        'string': 'VARCHAR(500)',
        'int4': 'INTEGER',
        'int': 'INTEGER',
        'integer': 'INTEGER',
        'varchar': 'VARCHAR(500)',
        'text': 'TEXT',
        'decimal': 'DECIMAL(18,2)',
        'numeric': 'NUMERIC(18,2)',
        'float': 'FLOAT',
        'double': 'DOUBLE PRECISION',
        'date': 'DATE',
        'datetime': 'TIMESTAMP',
        'timestamp': 'TIMESTAMP',
        'boolean': 'BOOLEAN',
        'bool': 'BOOLEAN'
    }
    
    # 处理带长度的varchar类型
    if original_type.lower().startswith('varchar('):
        return original_type.upper()
    
    # 处理decimal/numeric类型
    if original_type.lower().startswith(('decimal(', 'numeric(')):
        return original_type.upper()
    
    return type_mapping.get(original_type.lower(), 'VARCHAR(500)')

def format_column_name(column_name):
    """
    格式化字段名，处理特殊字符和以数字开头的情况
    
    Args:
        column_name (str): 原始字段名
        
    Returns:
        str: 格式化后的字段名
    """
    # 如果字段名以数字开头或包含特殊字符，使用双引号包围
    if (column_name[0].isdigit() or 
        not re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', column_name)):
        return f'"{column_name}"'
    
    return column_name

def read_columns_file_final(file_path):
    """
    最终版字段文件读取
    
    Args:
        file_path (str): columns.txt文件路径
        
    Returns:
        tuple: (字段列表, 错误行列表)
    """
    columns = []
    error_lines = []
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith('#'):  # 跳过空行和注释行
                    continue
                    
                parts = line.split('\t')
                if len(parts) >= 4:
                    table_name = parts[0].strip()
                    column_name = parts[1].strip()
                    data_type = parts[2].strip()
                    comment = parts[3].strip()
                    
                    # 基本验证：字段名不能为空
                    if column_name and data_type:
                        columns.append((table_name, column_name, data_type, comment))
                    else:
                        error_lines.append((line_num, f"字段名或数据类型为空: {line}"))
                else:
                    error_lines.append((line_num, f"格式不正确（需要4列）: {line}"))
                    
    except FileNotFoundError:
        print(f"错误：找不到文件 {file_path}")
        return [], []
    except Exception as e:
        print(f"错误：读取文件时发生异常 {e}")
        return [], []
        
    return columns, error_lines

def generate_final_alter_statements(columns, target_table_name, full_table_name):
    """
    生成最终版ALTER TABLE语句
    
    Args:
        columns (list): 字段信息列表
        target_table_name (str): 目标表名
        full_table_name (str): 完整表名
        
    Returns:
        list: SQL语句列表
    """
    statements = []
    
    # 过滤出属于目标表的字段
    target_columns = [col for col in columns if col[0] == target_table_name]
    
    if not target_columns:
        print(f"警告：没有找到表 {target_table_name} 的字段定义")
        return statements
    
    print(f"找到 {len(target_columns)} 个字段需要添加到表 {target_table_name}")
    
    # 添加文件头
    statements.extend([
        "-- ======================================================",
        "-- 高斯数据库表字段添加脚本（最终版）",
        "-- 自动生成的ALTER TABLE语句",
        f"-- 目标表: {full_table_name}",
        f"-- 字段总数: {len(target_columns)}",
        f"-- 生成时间: {__import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "-- ======================================================",
        "",
        "-- 执行说明：",
        "-- 1. 本脚本包含事务控制，确保操作的原子性",
        "-- 2. 如果某个字段已存在，整个事务会回滚",
        "-- 3. 建议先在测试环境执行验证",
        "-- 4. 执行前请确保有足够的权限和磁盘空间",
        "",
        "-- 开始事务",
        "BEGIN;",
        "",
        "-- 设置客户端编码",
        "SET client_encoding = 'UTF8';",
        ""
    ])
    
    # 生成ALTER TABLE语句
    statements.append("-- 添加字段")
    for i, (table_name, column_name, data_type, comment) in enumerate(target_columns, 1):
        # 格式化字段名
        formatted_column_name = format_column_name(column_name)
        
        # 映射数据类型
        mapped_type = map_data_type_final(data_type)
        
        # 生成ALTER TABLE语句
        alter_statement = f"ALTER TABLE {full_table_name} ADD COLUMN {formatted_column_name} {mapped_type};"
        statements.append(alter_statement)
        
        # 每100个字段添加一个进度注释
        if i % 100 == 0:
            statements.append(f"-- 已添加 {i}/{len(target_columns)} 个字段")
            statements.append("")
    
    statements.extend(["", "-- 添加字段注释"])
    
    # 生成注释语句
    for i, (table_name, column_name, data_type, comment) in enumerate(target_columns, 1):
        # 格式化字段名
        formatted_column_name = format_column_name(column_name)
        
        # 转义单引号
        escaped_comment = comment.replace("'", "''")
        
        # 生成注释语句
        comment_statement = f"COMMENT ON COLUMN {full_table_name}.{formatted_column_name} IS '{escaped_comment}';"
        statements.append(comment_statement)
        
        # 每100个注释添加一个进度注释
        if i % 100 == 0:
            statements.append(f"-- 已添加 {i}/{len(target_columns)} 个字段注释")
            statements.append("")
    
    # 添加文件尾
    statements.extend([
        "",
        "-- 提交事务",
        "COMMIT;",
        "",
        "-- 验证结果",
        f"SELECT COUNT(*) as column_count FROM information_schema.columns WHERE table_name = '{target_table_name.split('.')[-1]}' AND table_schema = '{full_table_name.split('.')[0]}';",
        "",
        "-- 脚本执行完成",
        f"-- 总共添加了 {len(target_columns)} 个字段",
        "-- 如果执行过程中出现错误，事务会自动回滚"
    ])
    
    return statements

def save_final_statements(statements, output_file, columns, error_lines):
    """
    保存最终版SQL语句
    
    Args:
        statements (list): SQL语句列表
        output_file (str): 输出文件路径
        columns (list): 字段列表
        error_lines (list): 错误行列表
    """
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            for statement in statements:
                f.write(statement + "\n")
                
        print(f"SQL语句已保存到文件: {output_file}")
        
        # 生成执行指南
        guide_file = output_file.replace('.sql', '_execution_guide.txt')
        with open(guide_file, 'w', encoding='utf-8') as f:
            f.write("高斯数据库字段添加脚本执行指南\n")
            f.write("=" * 50 + "\n\n")
            
            f.write("1. 执行前检查清单:\n")
            f.write("   □ 确认目标数据库连接正常\n")
            f.write("   □ 确认有ALTER TABLE权限\n")
            f.write("   □ 确认磁盘空间充足\n")
            f.write("   □ 已在测试环境验证\n")
            f.write("   □ 已备份相关数据\n\n")
            
            f.write("2. 执行命令示例:\n")
            f.write(f"   gsql -d database_name -f {os.path.basename(output_file)}\n\n")
            
            f.write("3. 脚本统计信息:\n")
            f.write(f"   总字段数: {len(columns)}\n")
            f.write(f"   有效字段数: {len([c for c in columns if c[0] == 'exec_imp_tag_bscom_bkupb_intg_upf_basic_info_np'])}\n")
            f.write(f"   错误行数: {len(error_lines)}\n\n")
            
            if error_lines:
                f.write("4. 处理错误详情:\n")
                for line_num, error in error_lines:
                    f.write(f"   第{line_num}行: {error}\n")
                f.write("\n")
            
            f.write("5. 注意事项:\n")
            f.write("   - 脚本使用事务，失败会自动回滚\n")
            f.write("   - 以数字开头的字段名会自动加双引号\n")
            f.write("   - 执行时间可能较长，请耐心等待\n")
            f.write("   - 建议在业务低峰期执行\n")
        
        print(f"执行指南已保存到: {guide_file}")
        
    except Exception as e:
        print(f"错误：保存文件时发生异常 {e}")

def main():
    """
    主函数
    """
    # 文件路径配置
    current_dir = os.path.dirname(os.path.abspath(__file__))
    columns_file = os.path.join(current_dir, 'columns.txt')
    output_file = os.path.join(current_dir, 'final_alter_table_statements.sql')
    
    # 目标表名配置
    target_table_name = 'exec_imp_tag_bscom_bkupb_intg_upf_basic_info_np'
    full_table_name = 'market.exec_imp_tag_bscom_bkupb_intg_upf_basic_info_np'
    
    print("高斯数据库表字段添加脚本（最终版）")
    print("=" * 50)
    print(f"输入文件: {columns_file}")
    print(f"输出文件: {output_file}")
    print(f"目标表名: {target_table_name}")
    print(f"完整表名: {full_table_name}")
    print("-" * 50)
    
    # 读取字段定义
    columns, error_lines = read_columns_file_final(columns_file)
    
    if not columns:
        print("没有读取到任何有效字段定义，程序退出")
        return
    
    print(f"总共读取到 {len(columns)} 个字段定义")
    if error_lines:
        print(f"发现 {len(error_lines)} 个错误行（已跳过）")
    
    # 统计数据类型分布
    type_stats = defaultdict(int)
    target_columns = [col for col in columns if col[0] == target_table_name]
    
    for _, _, data_type, _ in target_columns:
        mapped_type = map_data_type_final(data_type)
        type_stats[mapped_type] += 1
    
    print("\n数据类型分布:")
    for data_type, count in sorted(type_stats.items()):
        print(f"  {data_type}: {count} 个字段")
    
    # 生成ALTER TABLE语句
    statements = generate_final_alter_statements(columns, target_table_name, full_table_name)
    
    if not statements:
        print("没有生成任何SQL语句，程序退出")
        return
    
    sql_count = len([s for s in statements if s.strip() and not s.strip().startswith('--')])
    print(f"\n生成了 {sql_count} 条SQL语句")
    
    # 保存到文件
    save_final_statements(statements, output_file, columns, error_lines)
    
    print("-" * 50)
    print("✅ 脚本生成完成！")
    print("\n📋 下一步操作:")
    print("1. 查看生成的SQL文件")
    print("2. 阅读执行指南")
    print("3. 在测试环境验证")
    print("4. 在生产环境执行")

if __name__ == '__main__':
    main()